#! -*- coding:utf-8 -*-
'''pip install  gensim'''
from gensim import corpora, models, similarities
import logging
import jieba.analyse
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
#documents = ["Shipment of gold damaged in a fire","Delivery of silver arrived in a silver truck","Shipment of gold arrived in a truck"]

documents = []
for uid, line in enumerate(open("train.txt")):
    seg_list = jieba.cut(line)
    jiebastr=' '.join(seg_list)
    print jiebastr
    documents.append(jiebastr)
#print documents

texts = [[word for word in document.lower().split()] for document in documents]
print texts

dictionary = corpora.Dictionary(texts)
print dictionary
print "dictionary.token2id:",dictionary.token2id

corpus = [dictionary.doc2bow(text) for text in texts]
print "corpus:",corpus

tfidf = models.TfidfModel(corpus)
corpus_tfidf = tfidf[corpus]
for doc in corpus_tfidf:
    print "tfidf doc:",doc

#print tfidf.dfs
#print tfidf.idfs

lsi = models.LsiModel(corpus_tfidf, id2word=dictionary, num_topics=3)
lsi.print_topics(2)
lsi.save("lsi.mm")
corpus_lsi = lsi[corpus_tfidf]
for doc in corpus_lsi:
    print "corpus_lsi doc:",doc

index = similarities.MatrixSimilarity(lsi[corpus])

query = "hh"
for uid, line in enumerate(open("in.txt")):
    query=line
print "query:",query
query_seg_list = jieba.cut(query)
query_jiebastr=' '.join(query_seg_list)
query_bow = dictionary.doc2bow(query_jiebastr.split())
print query_bow
query_lsi = lsi[query_bow]
print "query_lsi:",query_lsi

sims = index[query_lsi]
print "sims index:",list(enumerate(sims))

sort_sims = sorted(enumerate(sims), key=lambda item: -item[1])
print "sort_sims:",sort_sims

print "sim doc:",documents[sort_sims[0][0]],"similar:",sort_sims[0][1]